#!/usr/bin/python
# -*- coding: us-ascii -*-
# vim:ts=4:sw=4:softtabstop=0:smarttab

"""
Test socket module
------------------

Test the new socket module.

"""
from __future__ import print_function
from __future__ import absolute_import

from pycopia import proctools
from pycopia.QA import core

USESOCKET = "pycopia"
#USESOCKET = "oldpycopia"
#USESOCKET = "stock"

if USESOCKET == "stock":
    import socket
    from pycopia.aid import systemcall

    class SafeSocket(socket._realsocket):
        """A socket protected from interrupted system calls."""
        accept = systemcall(socket._realsocket.accept)
        recv = systemcall(socket._realsocket.recv)
        send = systemcall(socket._realsocket.send)
        sendall = systemcall(socket._realsocket.sendall)
        connect = systemcall(socket._realsocket.connect)
        listen = systemcall(socket._realsocket.listen)
        bind = systemcall(socket._realsocket.bind)


    def connect_inet(host, port, socktype, sobject=SafeSocket):
        args = socket.getaddrinfo(str(host), int(port), socket.AF_INET, socktype)
        for family, socktype, proto, canonname, sockaddr in args:
            try:
                s = sobject(family, socktype, proto)
                s.connect(sockaddr)
            except:
                continue
            else:
                return s
        raise

    def connect_tcp(host, port, sobject=SafeSocket):
        return connect_inet(host, port, socket.SOCK_STREAM, sobject)

    def tcp_listener(addr, num=5, sobject=SafeSocket):
        s = sobject(socket.AF_INET, socket.SOCK_STREAM)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        s.bind(addr)
        s.listen(num)
        return s

    socket.connect_inet = connect_inet
    socket.connect_tcp = connect_tcp
    socket.tcp_listener = tcp_listener

elif USESOCKET == "pycopia":
    from pycopia import socket

elif USESOCKET == "oldpycopia":
    from pycopia import oldsocket as socket



class StartEchoServer(core.Test):
    """
    Purpose
    +++++++

    Start a test echo server.

    Pass Criteria
    +++++++++++++

    Start without errors.

    Start Condition
    +++++++++++++++

    Nothing special.

    End Condition
    +++++++++++++

    Test server is running.

    Reference
    +++++++++

    TBD

    Prerequisites
    +++++++++++++

    None

    Procedure
    +++++++++

    1. Start the server with the TestWorker, as a subprocess.
    2. Store the subprocess server in the config. 
    """

    def execute(self):
        self.info("Starting server.")
        self.config.server = proctools.submethod(_run_server, args=(8999, bool(self.config.flags.DEBUG)))
        self.sleep(2)
        return self.passed("Started server")


def _run_server(port, debug):

    class LineEchoServer(object):
        def __init__(self, port, host=""):
            self._sock = socket.tcp_listener((host, port), 1)

        def listen(self):
            conn, addr = self._sock.accept()
            fo = conn.makefile("w+b", 0)
            while 1:
                data = fo.readline()
                if not data: 
                    break
                if data.startswith("BYE"):
                    break
                fo.write(data)
            fo.close()
            conn.close()

        def run(self):
            try:
                try:
                    while 1:
                        self.listen()
                except KeyboardInterrupt:
                    return
            finally:
                self._sock.close()

    class StreamEchoServer(object):
        def __init__(self, port, host=""):
            self._sock = socket.tcp_listener((host, port), 1)

        def listen(self):
            conn, addr = self._sock.accept()
            while 1:
                data = conn.recv(4096)
                if not data: 
                    break
                if data.startswith("BYE"):
                    break
                conn.send(data)
            conn.close()

        def run(self):
            try:
                try:
                    while 1:
                        self.listen()
                except KeyboardInterrupt:
                    return
            finally:
                self._sock.close()

    #srv = StreamEchoServer(port, "localhost")
    srv = LineEchoServer(port, "localhost")
    srv.run()


class MakefileTest(core.Test):
    """
    Purpose
    +++++++

    Test that the file-like object returned from makefile method works properly.

    Pass Criteria
    +++++++++++++

    TBD

    Start Condition
    +++++++++++++++

    echo server running on port 8999

    End Condition
    +++++++++++++

    No change.

    Reference
    +++++++++

    pycopia/aid/socket.py

    Prerequisites
    +++++++++++++

    None

    Procedure
    +++++++++

    1. Step 1 ...
    2. Step 2 ...
    3. Step 3 ...

    """
    PREREQUISITES = ["StartEchoServer"]

    def execute(self):
        s = socket.connect_tcp("localhost", 8999)
        f = s.makefile("w+b", 0)
        LINE = "hello from test\n"
        f.write(LINE)
        rline = f.readline()
        f.write("BYE\n")
        f.close()
        s.close()
        if rline == LINE:
            self.info("got {!r}".format(rline))
            return self.passed("Wrote to file-like object.")
        else:
            self.diagnostic("got {!r}".format(rline))
            return self.failed("Didn't get expected line back.")


class StopServer(core.Test):
    """
    Purpose
    +++++++

    Stop the server. This is a helper test.

    Pass Criteria
    +++++++++++++

    Server stops cleanly.

    Start Condition
    +++++++++++++++

    Server is running.

    End Condition
    +++++++++++++

    No server is running.

    Reference
    +++++++++

    None

    Prerequisites
    +++++++++++++

    StartEchoServer

    Procedure
    +++++++++

    1. Kill the server subprocess.
    2. Wait for exit status.

    """
    PREREQUISITES = ["StartEchoServer"]

    def execute(self):
        try:
            srv = self.config.server
        except AttributeError:
            return self.incomplete("No server")
        del self.config.server
        srv.kill()
        es = srv.wait()
        if es:
            self.info(es)
        else:
            self.diagnostic(es)
        return self.passed("Killed server.")




def get_suite(config):
    suite = core.TestSuite(config, name="SocketSuite")
    suite.add_test(StartEchoServer)
    suite.add_test(MakefileTest)
    suite.add_test(StopServer)
    return suite

def run(config):
    suite = get_suite(config)
    suite.run()

